from re import T
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.optim as optim
import torch.utils.data
import torch.utils.data.distributed
import torchvision.datasets as datasets
import torchvision.transforms as transforms
from sklearn.metrics import classification_report
from timm.data.mixup import Mixup
from timm.loss import SoftTargetCrossEntropy
from timm.models import swin_small_patch4_window7_224
from torchtoolbox.transform import Cutout
import os

# Set global parameters
GPU_ID = '0'
os.environ["CUDA_VISIBLE_DEVICES"]= GPU_ID
DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
BATCH_SIZE = 64
DATA_PATH = './data4train'
CKPT_PATH = ''

# preprocessing
transform_test = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# Loading DataLoader
dataset_test = datasets.ImageFolder(DATA_PATH + "/val", transform=transform_test)
test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=BATCH_SIZE, shuffle=False)

# Loading model
model_ft = torch.load(CKPT_PATH)
model_ft.eval()
model_ft.to(DEVICE)
criterion_val = torch.nn.CrossEntropyLoss()

# define val function
ACC = 0
def val(model, device, test_loader):
    global ACC
    model.eval()
    test_loss = 0
    correct = 0
    total_num = len(test_loader.dataset)
    print(total_num, len(test_loader))
    val_list = []
    pred_list = []
    with torch.no_grad():
        for data, target in test_loader:
            for t in target:
                val_list.append(t.data.item())
            data, target = data.to(device), target.to(device)
            output = model(data)
            loss = criterion_val(output, target)
            _, pred = torch.max(output.data, 1)
            for p in pred:
                pred_list.append(p.data.item())
            correct += torch.sum(pred == target)
            print_loss = loss.data.item()
            test_loss += print_loss
        correct = correct.data.item()
        acc = correct / total_num
        avgloss = test_loss / len(test_loader)
        print('\nVal set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
            avgloss, correct, len(test_loader.dataset), 100 * acc))
    return val_list, pred_list

# main
val_list, pred_list = val(model_ft, DEVICE, test_loader)
print(classification_report(val_list, pred_list, target_names=dataset_test.class_to_idx))
print('Real Target : ')
print(val_list)
print('Predict Target :')
print(pred_list)